__author__ = 'shawhen'

import ctypes

from . import WinNT
from . import Kernel32
from . import WinDef
from . import BaseTsd

#
# The th32ProcessID argument is only used if TH32CS_SNAPHEAPLIST or
# TH32CS_SNAPMODULE is specified. th32ProcessID == 0 means the current
# process.
#
# NOTE that all of the snapshots are global except for the heap and module
#      lists which are process specific. To enumerate the heap or module
#      state for all WIN32 processes call with TH32CS_SNAPALL and the
#      current process. Then for each process in the TH32CS_SNAPPROCESS
#      list that isn't the current process, do a call with just
#      TH32CS_SNAPHEAPLIST and/or TH32CS_SNAPMODULE.
#
# dwFlags
#
TH32CS_SNAPHEAPLIST = 0x00000001
TH32CS_SNAPPROCESS  = 0x00000002
TH32CS_SNAPTHREAD   = 0x00000004
TH32CS_SNAPMODULE   = 0x00000008
TH32CS_SNAPMODULE32 = 0x00000010
TH32CS_SNAPALL      = (TH32CS_SNAPHEAPLIST | TH32CS_SNAPPROCESS | TH32CS_SNAPTHREAD | TH32CS_SNAPMODULE)
TH32CS_INHERIT      = 0x80000000
#
# Use CloseHandle to destroy the snapshot
#
def CreateToolhelp32Snapshot(dwFlags, dwTh32ProcessID):
    """
    :return HANDLE
    """
    dwFlags_c = WinDef.DWORD(dwFlags)
    dwTh32ProcessID_c = WinDef.DWORD(dwTh32ProcessID)
    return Kernel32.Kernel32_dll.CreateToolhelp32Snapshot(dwFlags_c, dwTh32ProcessID_c)


class PROCESSENTRY32W(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("dwSize", WinDef.DWORD),
        ("cntUsage", WinDef.DWORD),
        ("th32ProcessID", WinDef.DWORD),
        ("th32DefaultHeapID", BaseTsd.ULONG_PTR),
        ("th32ModuleID", WinDef.DWORD),
        ("cntThreads", WinDef.DWORD),
        ("th32ParentProcessID", WinDef.DWORD),
        ("pcPriClassBase", WinNT.LONG),
        ("dwFlags", WinDef.DWORD),
        ("szExeFile", WinNT.WCAHR*WinDef.MAX_PATH)
    ]

def Process32First(hSnapshot, pe):
    """
    :param hSnapshot 快照句柄
    :param pe 进程信息结构体（不是指针哟）
    :return BOOL
    """
    return Kernel32.Kernel32_dll.Process32FirstW(hSnapshot, ctypes.byref(pe))


def Process32Next(hSnapshot, pe):
    """
    同上
    """
    return Kernel32.Kernel32_dll.Process32NextW(hSnapshot, ctypes.byref(pe))